from scipy.stats import t, f
import numpy as np
import math


# 1. Function for One-Sample t-Test


def one_sample_t_test(alpha, sides, sample_mean, population_mean, std_dev, power):
    """
    Calculate sample size for a one-sample t-test.

    Parameters:
    - alpha: significance level (e.g., 0.05 for 5%)
    - sides: 1 for one-sided test, 2 for two-sided test
    - sample_mean: mean of the sample
    - population_mean: mean of the population (null hypothesis mean)
    - std_dev: standard deviation of the sample
    - power: desired power (e.g., 85 for 85%)

    Returns:
    - Sample size required for the one-sample t-test
    """
    n = 2  # Initial sample size

    while True:
        df = n - 1  # Degrees of freedom
        effect_size = abs(sample_mean - population_mean) / std_dev  # Effect size
        ncp = math.sqrt(n) * effect_size  # Non-centrality parameter

        if sides == 1:  # One-sided test
            t_crit = t.ppf(1 - alpha, df)
            achieved_power = 1 - t.cdf(t_crit, df, ncp)
        elif sides == 2:  # Two-sided test
            t_crit_1 = t.ppf(1 - alpha / 2, df)
            t_crit_2 = t.ppf(alpha / 2, df)
            achieved_power = 1 - t.cdf(t_crit_1, df, ncp) + t.cdf(t_crit_2, df, ncp)

        if achieved_power >= power / 100:
            return math.ceil(n)

        n += 0.01


# 3. Function for One-Way Repeated Measures ANOVA

def repeated_measures_anova(alpha, num_levels, std_dev, correlation, power, means):
    """
    Calculate sample size for a one-way repeated measures ANOVA.

    Parameters:
    - alpha: significance level (e.g., 0.05 for 5%)
    - num_levels: number of repeated measurement levels
    - std_dev: standard deviation of the sample at each level
    - correlation: correlation between repeated measurements
    - power: desired power (e.g., 85 for 85%)
    - means: list of means for each level

    Returns:
    - Sample size required for the repeated measures ANOVA
    """
    n = 2  # Initial sample size
    mean_variance = np.var(means, ddof=1)  # Variance of the means across levels
    effect_size = mean_variance / (std_dev ** 2 * (1 - correlation))  # Effect size

    while True:
        df1 = num_levels - 1  # Between-level degrees of freedom
        df2 = (n - 1) * (num_levels - 1)  # Within-level degrees of freedom
        ncp = n * num_levels * effect_size  # Non-centrality parameter

        f_crit = f.ppf(1 - alpha, df1, df2)  # Critical F-value
        achieved_power = 1 - f.cdf(f_crit, df1, df2, ncp)  # Achieved power

        if achieved_power >= power / 100:
            return math.ceil(n)

        n += 0.01


# Example 1: One-Sample t-Test
example_1 = one_sample_t_test(0.05, 2, 130.83, 140, 25.74, 85)


print("Example 1:", example_1, "samples needed for One-Sample t-Test")

